package com.baidb.gateway.ribbon.filter;

import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR;

import com.baidb.gateway.properties.GrayProperties;
import java.net.URI;
import java.util.HashMap;
import java.util.Map;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
import org.springframework.cloud.gateway.config.LoadBalancerProperties;
import org.springframework.cloud.gateway.filter.LoadBalancerClientFilter;
import org.springframework.cloud.netflix.ribbon.RibbonLoadBalancerClient;
import org.springframework.http.HttpHeaders;
import org.springframework.web.server.ServerWebExchange;

/**
 * @author baidb
 * @date 2021/6/17.
 */
@SuppressWarnings("deprecation")
public class GrayLoadBalancerClientFilter extends LoadBalancerClientFilter {

    private GrayProperties grayProperties;

    public GrayLoadBalancerClientFilter(LoadBalancerClient loadBalancer,
        LoadBalancerProperties properties, GrayProperties grayProperties) {
        super(loadBalancer, properties);
        this.grayProperties = grayProperties;
    }

    @Override
    protected ServiceInstance choose(ServerWebExchange exchange) {
        if (grayProperties.isEnableGrayRelease() && this.loadBalancer instanceof RibbonLoadBalancerClient) {
            RibbonLoadBalancerClient client = (RibbonLoadBalancerClient)this.loadBalancer;
            HttpHeaders headers = exchange.getRequest().getHeaders();
            // 实现想要的路由规则逻辑
            String version = headers.getFirst("version");
            Map<String, String> hint = new HashMap<>();
            hint.put("version", version);
            String serviceId = ((URI)exchange.getAttribute(GATEWAY_REQUEST_URL_ATTR)).getHost();
            return client.choose(serviceId, hint);
        }
        return super.choose(exchange);
    }

}
